from __future__ import annotations
from pathlib import Path
import subprocess
from Bio.PDB.PDBParser import PDBParser
from typing import Iterator, TYPE_CHECKING, NamedTuple
import numpy as np
import pandas as pd

if TYPE_CHECKING:
    from Bio.PDB.Residue import Residue
    from Bio.PDB.Atom import Atom
    from Bio.PDB.Structure import Structure

dockin = '''
Receptor
{name}.pdb

RMSD
1.0

Binding pocket
{x_min} {x_max}
{y_min} {y_max}
{z_min} {z_max}

Number of binding poses
20

Ligands list
ligands.list

END
'''

class BindingPocket(NamedTuple):
    x_min: float
    x_max: float
    y_min: float
    y_max: float
    z_min: float
    z_max: float


def get_residue_coordinate(residue: Residue) -> list[np.ndarray]:
    coordinates = []
    for atom in residue.get_atoms():
        if atom.element != "H":
            coordinates.append(atom.coord)
    return coordinates

def get_pocket_coordiate(residue_ids: str, pdb_structure: PDBParser, extented_range: float = 6.0) -> BindingPocket:
    coordinates = []
    for residue_id in residue_ids.split(' '):
        chain_name, id_name = residue_id.split('_')
        coordinates.extend(get_residue_coordinate(pdb_structure[chain_name][int(id_name)]))
    coordinate = np.row_stack(coordinates)
    min_coordinate = coordinate.min(axis=0)
    min_coordinate = min_coordinate - extented_range
    max_coordinate = coordinate.max(axis=0)
    max_coordinate = max_coordinate + extented_range
    return BindingPocket(
        x_min = min_coordinate[0],
        x_max = max_coordinate[0],
        y_min = min_coordinate[1],
        y_max = max_coordinate[1],
        z_min = min_coordinate[2],
        z_max = max_coordinate[2],
        )

def create_dockin_file(output_dir: Path, pdb_name: str,
                       binding_pocket_rank: str, binding_pocket: BindingPocket) -> None:
    with open(output_dir / f'{pdb_name}_{binding_pocket_rank}.in', 'w+', encoding='utf-8', newline='\n') as f:
        f.write(dockin.format(
            name = f'targets/pdb{pdb_name}',
            x_min = binding_pocket.x_min,
            x_max = binding_pocket.x_max,
            y_min = binding_pocket.y_min,
            y_max = binding_pocket.y_max,
            z_min = binding_pocket.z_min,
            z_max = binding_pocket.z_max,
        ))

def iterate_pocket_file(pocket_dir: Path, pdb_dir: Path, output_dir: Path,
                        pdb_parser: PDBParser) -> None:
    for filename in pocket_dir.iterdir():
        if filename.suffix == '.csv':
            if filename.stem.split('.')[1] == 'pdb_predictions':
                pdb_name = filename.stem.split('.')[0][3:]
                pdb_structure = pdb_parser.get_structure(pdb_name, pdb_dir/f'pdb{pdb_name}.pdb')[0]
                predict_data = pd.read_csv(filename, skipinitialspace=True)
                for _, row in predict_data.iterrows():
                    rank = row['rank']
                    residue_ids = row['residue_ids']
                    binding_pocket = get_pocket_coordiate(residue_ids, pdb_structure)
                    create_dockin_file(output_dir, pdb_name, rank, binding_pocket)

if __name__ == '__main__':
    # 重要，需要替换为真实的p2rank蛋白结合位点预测结果输出的路径
    pocket_dir = Path('重要，需要替换为真实的p2rank蛋白结合位点预测结果输出的路径')
    output_dir = Path('.')
    pdb_dir = output_dir / 'targets'

    pdb_parser = PDBParser(PERMISSIVE=True, QUIET=True)

    iterate_pocket_file(pocket_dir, pdb_dir, output_dir, pdb_parser)